import paddle.fluid as fluid

from helm.static.models.layers import Conv2d, Pool2d, GlobalAvgPool, BN, Linear, Act, Layer, Sequential
from helm.static.models.modules import SplAtConv2d

__all__ = [
    "PyramidNeSt"
]


class PadChannel(Layer):
    def __init__(self, channels):
        super().__init__()
        assert channels >= 0
        self.c = channels

    def forward(self, x):
        return fluid.layers.pad(x, [0, 0, 0, self.c, 0, 0, 0, 0])

    def extra_repr(self):
        return "padding=%d" % self.c


class Shortcut(Layer):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        layers = []
        if stride == 2:
            layers.append(("pool", Pool2d(2, 2)))
        if in_channels != out_channels:
            layers.append(("pad", PadChannel(out_channels - in_channels)))
        self.shortcut = Sequential(*layers)

    def forward(self, x):
        return self.shortcut(x)


class Bottleneck(Layer):
    expansion = 1

    def __init__(self, in_channels, channels, groups, stride=1, radix=2):
        super().__init__()
        out_channels = channels * self.expansion
        branch1 = [
            ("bn0", BN(in_channels)),
            ("conv1", Conv2d(in_channels, channels, kernel_size=1, bn=True, act='default')),
            *([("pool", Pool2d(3, 2))] if stride != 1 else []),
            ("conv2",
             SplAtConv2d(channels, channels, kernel_size=3, groups=groups, radix=radix)
             if radix != 0 else Conv2d(channels, channels, kernel_size=3, groups=groups, bn=True, act='default')),
            ("conv3", Conv2d(channels, out_channels, kernel_size=1, bn=True)),
        ]
        self.branch1 = Sequential(*branch1)
        self.branch2 = Shortcut(in_channels, out_channels, stride)

    def forward(self, x):
        return self.branch1(x) + self.branch2(x)


def round_channels(channels, divisor=8, min_depth=None):
    min_depth = min_depth or divisor
    new_channels = max(min_depth, int(channels + divisor / 2) // divisor * divisor)
    if new_channels < 0.9 * channels:
        new_channels += divisor
    return int(new_channels)


class PyramidNeSt(Layer):
    def __init__(self, start_channels, widening_fractor, depth, groups, radix, num_classes=10):
        super().__init__()

        num_layers = [(depth - 2) // 9] * 3

        strides = [1, 2, 2]

        self.add_channel = widening_fractor / sum(num_layers)
        self.in_channels = start_channels
        self.channels = start_channels

        layers = [("init_block", Conv2d(3, start_channels, kernel_size=3, bn=True))]

        for i, (n, s) in enumerate(zip(num_layers, strides)):
            layers.append(
                ("stage%d" % (i + 1), self._make_layer(n, groups, stride=s, radix=radix))
            )

        layers.append(("post_activ", Sequential(
            ("bn", BN(self.in_channels)),
            ("act", Act()),
        )))
        self.features = Sequential(*layers)
        assert (start_channels + widening_fractor) * Bottleneck.expansion == self.in_channels
        self.final_pool = GlobalAvgPool()
        self.output = Linear(self.in_channels, num_classes)

    def _make_layer(self, num_layers, groups, stride, radix):
        self.channels = self.channels + self.add_channel
        layers = [("unit1", Bottleneck(self.in_channels, round_channels(self.channels, groups),
                                       groups, stride=stride, radix=radix))]
        self.in_channels = round_channels(self.channels, groups) * Bottleneck.expansion
        for i in range(1, num_layers):
            self.channels = self.channels + self.add_channel
            layers.append(("unit%d" % (i + 1), Bottleneck(self.in_channels, round_channels(self.channels, groups),
                                                          groups, radix=radix)))
            self.in_channels = round_channels(self.channels, groups) * Bottleneck.expansion
        return Sequential(*layers)

    def forward(self, x):
        x = self.features(x)
        x = self.final_pool(x)
        x = self.output(x)
        return x
